#!/usr/bin/env python
# -*- encoding: utf-8 -*-
'''
@File     : client.py
@Project  : pipecoco
@Date     : 2021/8/13
@Author   : Zhang Jinyang
@Contact  : zhang-jy@sjtu.edu.cn
'''

from comm.pipecoco_client import pipecoco_client
from net.alexnet.loader import load_alexnet
from net.vgg16.loader import load_vgg_small, load_vgg_large
from net.resnet.loader import load_resnet34
from net.config import input_shapes,divided_points
from net.datasets import create_dataset_cifar10

# 创建PipeCoCo客户端示例
class client_test(pipecoco_client):

    def __init__(self, host, record_path):
        pipecoco_client.__init__(self, host, record_path)

    def load_net(self, model_name):

        if model_name == 'alexnet':
            network = load_alexnet('./models/alexnet.ckpt')

        elif model_name == 'vgg_small':
            network = load_vgg_small('./models/vgg_small.ckpt')

        elif model_name == 'vgg_large':
            network = load_vgg_large('./models/vgg_large.ckpt')

        elif model_name == ('resnet34'):
            network = load_resnet34('./models/resnet34.ckpt')

        return network

if __name__=='__main__':

    # 加载数据集
    batch_size = 20
    model_name = 'alexnet'
    input_shape = [3, 227, 227]
    divided_point = 2
    ds_eval = create_dataset_cifar10('./cifar10_dataset/train', input_shape[1:], batch_size=batch_size).create_dict_iterator()

    # 常见客户端
    host = 'localhost:5000'
    client = client_test(host,'record/pi2server_batch_size=20')
    input = next(ds_eval)['image']
    client.create_model(model_name, input_shape, divided_point=divided_point, input=input, optim=False, batch_size = 20)
    client.optimize(model_name, input)
    
    while True:
        input = next(ds_eval)['image']
        result = client.infer(input, model_name)
        print("result: ", result)
